#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# 首先导入修复脚本，解决三个主要问题：
# 1. HuggingFace transformers 的 chat_template 设置问题
# 2. BERTScore 的 meta tensor 设备迁移问题
# 3. METEOR 的 NLTK 语料库兼容问题
import sys
import os
# 将脚本目录添加到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
# 导入修复脚本
from fix_model_initialization import apply_all_fixes
# 应用所有修复
apply_all_fixes()

import os
import json
import torch
import logging
from datetime import datetime

# 设置日志
torch.cuda.empty_cache()
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("patient_model_test.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 导入自定义模块
from data_processor import load_patient_data, prepare_context_data
from model_handler import load_model, generate_patient_response
from evaluator import calculate_all_metrics
from experiment_manager import run_experiments
from result_manager import save_results, generate_report
from config import DATA_DIR, MODEL_PATH, EXPERIMENT_CONFIG


def main():
    """主函数，执行整个患者模型测试流程"""
    logger.info("===== 开始执行患者模型测试流程 =====")
    start_time = datetime.now()
    
    try:
        # 1. 加载患者数据
        logger.info(f"加载患者数据从目录: {DATA_DIR}")
        patient_data_list = load_patient_data(DATA_DIR)
        logger.info(f"成功加载 {len(patient_data_list)} 个患者数据文件")
        
        # 2. 加载模型
        logger.info(f"加载模型: {MODEL_PATH}")
        tokenizer, model = load_model(MODEL_PATH)
        
        # 3. 执行实验
        logger.info("开始执行实验")
        results = run_experiments(
            patient_data_list,
            tokenizer,
            model,
            EXPERIMENT_CONFIG
        )
        
        # 4. 保存结果
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        result_file = f"results_{timestamp}.json"
        logger.info(f"保存结果到文件: {result_file}")
        save_results(results, result_file)
        
        # 5. 生成报告
        report_file = f"report_{timestamp}.md"
        logger.info(f"生成实验报告: {report_file}")
        generate_report(results, report_file)
        
        end_time = datetime.now()
        execution_time = (end_time - start_time).total_seconds()
        logger.info(f"===== 患者模型测试流程执行完成，总耗时: {execution_time:.2f} 秒 =====")
        
    except Exception as e:
        logger.error(f"执行过程中发生错误: {str(e)}", exc_info=True)
    

if __name__ == "__main__":
    main()